{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d15dad13-9732-4e4c-bbd1-1a33545a4293",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f857fc7-d7fb-4b05-a242-de31fb1f086d",
   "metadata": {},
   "source": [
    "In this notebook, we'll go through the process of fine-tuning the LightGBM models in the `pdf-document-layout-analysis` service."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c96b645-eef0-47a2-8c4f-284cdc05e76d",
   "metadata": {},
   "source": [
    "But before doing that, let's start with some basic concepts and introduce modules and methods to make the process easier and cleaner."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1e5c19b-1920-4f2c-9994-943626cd8a58",
   "metadata": {},
   "source": [
    "To begin with, you should first ensure that `Poppler` is installed on your system. We will use it to process PDFs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5f198930-caf1-4cb4-bb1e-8ca063ad8587",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pdftohtml is already installed.\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "\n",
    "if ! command -v pdftohtml &> /dev/null\n",
    "then\n",
    "    echo \"pdftohtml is not installed. Installing now...\"\n",
    "    sudo apt install pdftohtml\n",
    "else\n",
    "    echo \"pdftohtml is already installed.\"\n",
    "fi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d971faa-e9a8-47d6-8c02-66be6f3a3c6c",
   "metadata": {},
   "source": [
    "We use Poppler to convert PDFs to XMLs. To work with Poppler in Python, we have created `PdfFeatures` module, which can be found in `pdf_features/PdfFeatures.py`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f7ac5d42-fb70-4476-8e05-b159f18ae3dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pdf_features.PdfFeatures import PdfFeatures"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e45522eb-6879-472a-a822-64b38041ccc3",
   "metadata": {},
   "source": [
    "To open a PDF file with PdfFeatures module, simply write:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "e4ac53e5-b249-4dcd-beeb-e3009e17079b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Page-1\n",
      "Page-2\n"
     ]
    }
   ],
   "source": [
    "pdf_features: PdfFeatures = PdfFeatures.from_pdf_path(\"test_pdfs/regular.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7c6241-9016-4416-a53e-644145f9063a",
   "metadata": {},
   "source": [
    "When you open `pdf_features` like this, the XML file is saved in a temporary path and handled on the fly.\n",
    "\n",
    "If you want to save the XML file, you should provide a path where it can be saved:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "eb1056ee-2e45-4b12-b2bc-8d23553c2143",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Page-1\n",
      "Page-2\n"
     ]
    }
   ],
   "source": [
    "pdf_features: PdfFeatures = PdfFeatures.from_pdf_path(\"test_pdfs/regular.pdf\", \"test_pdfs/regular.xml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "703ec555-c3a5-4e7e-a6dd-886be67cb6de",
   "metadata": {},
   "source": [
    "Here is a part of the XML to illustrate what it looks like:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6fcebd-f91b-43fe-b2d6-b9956c3fd173",
   "metadata": {},
   "source": [
    "```\n",
    "<?xml version=\"1.0\" encoding=\"UTF-8\"?>\n",
    "<!DOCTYPE pdf2xml SYSTEM \"pdf2xml.dtd\">\n",
    "\n",
    "<pdf2xml producer=\"poppler\" version=\"23.04.0\">\n",
    "<page number=\"1\" position=\"absolute\" top=\"0\" left=\"0\" height=\"842\" width=\"595\">\n",
    "\t<fontspec id=\"0\" size=\"10\" family=\"JOIBEJ+Verdana\" color=\"#000000\"/>\n",
    "\t<fontspec id=\"1\" size=\"10\" family=\"JOIBGK+Verdana\" color=\"#000000\"/>\n",
    "<text top=\"106\" left=\"244\" width=\"111\" height=\"12\" font=\"0\"><b>RESOLUCIÓN DE LA </b></text>\n",
    "<text top=\"118\" left=\"157\" width=\"284\" height=\"12\" font=\"0\"><b>CORTE INTERAMERICANA DE DERECHOS HUMANOS </b></text>\n",
    "<text top=\"129\" left=\"227\" width=\"145\" height=\"12\" font=\"0\"><b>DEL 29 DE JULIO DE 1991 </b></text>\n",
    "<text top=\"141\" left=\"298\" width=\"3\" height=\"12\" font=\"0\"><b> </b></text>\n",
    "<text top=\"153\" left=\"298\" width=\"3\" height=\"12\" font=\"0\"><b> </b></text>\n",
    "<text top=\"165\" left=\"132\" width=\"334\" height=\"12\" font=\"0\"><b>MEDIDAS PROVISIONALES SOLICITADAS POR LA COMISIÓN </b></text>\n",
    "<text top=\"177\" left=\"177\" width=\"245\" height=\"12\" font=\"0\"><b>INTERAMERICANA DE DERECHOS HUMANOS </b></text>\n",
    "<text top=\"188\" left=\"225\" width=\"149\" height=\"12\" font=\"0\"><b>RESPECTO DE GUATEMALA </b></text>\n",
    "\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be01120-c4ce-4e09-bc10-64b1742c9b0b",
   "metadata": {},
   "source": [
    "When we convert PDFs to XMLs with Poppler, it creates `tokens`. These tokens are generally lines of text, but they can vary according to Poppler's heuristics and what has been extracted.  \n",
    "A token can be a single character, empty string, or an entire line. Every `<text>` item you see above is a `token`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00517165-bc84-4a6f-9a8b-91084cc603ab",
   "metadata": {},
   "source": [
    "The PdfFeatures module provides basic capabilities for working with PDF files. Here are some features of this module.  \n",
    "You don't have to memorize them, but they can be useful for future reference:\n",
    "\n",
    "- Every PdfFeatures instance has `pages` attribute. This attribute includes a list of `PdfPage` elements to work with each of the pages.\n",
    "- Every PdfPage element has attributes like `page_number`, `page_width`, `page_height` and `tokens`.\n",
    "- The `tokens` attribute includes a list of `PdfToken` elements to work with each of the tokens within that page.\n",
    "- Every PdfToken element has attributes like `content`, `bounding_box`, `token_type`, `page_number`.\n",
    "- The `content` attribute is, as the name implies, the string content of the given token.\n",
    "- The`bounding_box` attribute stores the position of the given token on the page.\n",
    "- `bounding_box` is a `Rectangle` element. For example, if you want to get the left coordinate of the token, you can do so by typing `token.bounding_box.left`. It will return an integer value.\n",
    "- `token_type` attribute is for keeping the type of the token. It's a `TokenType` element and you'll see more details about this one in the next sections.\n",
    "- Like PdfPage items, tokens also have a `page_number` attribute to indicate which page they are on. This is useful in some scenarios."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63a71904-0ad3-4fca-830a-402d9334614a",
   "metadata": {},
   "source": [
    "If you want to loop through the tokens of a file and check their contents you can use something like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "444d3778-c3f5-48fd-aa20-cfe1bf851aad",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[96mRESOLUCIÓN DE LA\u001b[0m \u001b[93m[Page: 1 || Coordinates: [244, 106, 355, 118]]\u001b[0m\n",
      "\u001b[96mCORTE INTERAMERICANA DE DERECHOS HUMANOS\u001b[0m \u001b[93m[Page: 1 || Coordinates: [157, 118, 441, 130]]\u001b[0m\n",
      "\u001b[96mDEL 29 DE JULIO DE 1991\u001b[0m \u001b[93m[Page: 1 || Coordinates: [227, 129, 372, 141]]\u001b[0m\n",
      "\u001b[96mMEDIDAS PROVISIONALES SOLICITADAS POR LA COMISIÓN\u001b[0m \u001b[93m[Page: 1 || Coordinates: [132, 165, 466, 177]]\u001b[0m\n",
      "\u001b[96mINTERAMERICANA DE DERECHOS HUMANOS\u001b[0m \u001b[93m[Page: 1 || Coordinates: [177, 177, 422, 189]]\u001b[0m\n",
      "\u001b[96mRESPECTO DE GUATEMALA\u001b[0m \u001b[93m[Page: 1 || Coordinates: [225, 188, 374, 200]]\u001b[0m\n",
      "\u001b[96mCASO CHUNIMA\u001b[0m \u001b[93m[Page: 1 || Coordinates: [254, 224, 344, 236]]\u001b[0m\n",
      "\u001b[96mLA CORTE INTERAMERICANA DE DERECHOS HUMANOS,\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 259, 393, 271]]\u001b[0m\n",
      "\u001b[96mVISTOS:\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 295, 137, 307]]\u001b[0m\n",
      "\u001b[96m1.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 318, 101, 330]]\u001b[0m\n",
      "\u001b[96mLa resolución del Presidente de la Corte Interamericana de Derechos Humanos\u001b[0m \u001b[93m[Page: 1 || Coordinates: [122, 318, 511, 330]]\u001b[0m\n",
      "\u001b[96mde 15 de julio de 1991, sobre medidas provisionales solicitadas por la Comisión\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 330, 514, 342]]\u001b[0m\n",
      "\u001b[96mInteramericana de Derechos Humanos respecto de Guatemala;\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 342, 401, 354]]\u001b[0m\n",
      "\u001b[96m2.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 366, 102, 378]]\u001b[0m\n",
      "\u001b[96mLa convocatoria a una audiencia pública para el día 29 de julio de 1991 a las\u001b[0m \u001b[93m[Page: 1 || Coordinates: [122, 366, 512, 378]]\u001b[0m\n",
      "\u001b[96m3:00 p.m., contenida en la resolución citada;\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 378, 312, 390]]\u001b[0m\n",
      "\u001b[96m3.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 401, 104, 413]]\u001b[0m\n",
      "\u001b[96mLos escritos de fechas 24 y 26 de este mes de julio presentados por el\u001b[0m \u001b[93m[Page: 1 || Coordinates: [122, 401, 514, 413]]\u001b[0m\n",
      "\u001b[96mGobierno de Guatemala en los cuales informa que, en atención a la resolución del\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 413, 513, 425]]\u001b[0m\n",
      "\u001b[96mPresidente, ha tomado medidas dirigidas a la protección de las personas\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 425, 518, 437]]\u001b[0m\n",
      "\u001b[96mmencionadas en esa resolución y solicita un aplazamiento de por lo menos 30 días de\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 437, 512, 449]]\u001b[0m\n",
      "\u001b[96mla audiencia convocada por el Presidente para hoy, a fin de contar con un plazo que\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 448, 512, 460]]\u001b[0m\n",
      "\u001b[96mle permita hacer una presentación adecuada ante la Corte.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 460, 380, 472]]\u001b[0m\n",
      "\u001b[96mCONSIDERANDO:\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 484, 189, 496]]\u001b[0m\n",
      "\u001b[96m1.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 508, 101, 520]]\u001b[0m\n",
      "\u001b[96mQue, en virtud del artículo 23.4 de su Reglamento, la Corte Interamericana de\u001b[0m \u001b[93m[Page: 1 || Coordinates: [122, 508, 511, 520]]\u001b[0m\n",
      "\u001b[96mDerechos Humanos debe pronunciarse sobre la resolución del Presidente del 15 de\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 519, 513, 531]]\u001b[0m\n",
      "\u001b[96mjulio de 1991;\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 531, 160, 543]]\u001b[0m\n",
      "\u001b[96m2.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 555, 104, 567]]\u001b[0m\n",
      "\u001b[96mQue, habida cuenta de que la Corte se encuentra reunida, debe también\u001b[0m \u001b[93m[Page: 1 || Coordinates: [122, 555, 514, 567]]\u001b[0m\n",
      "\u001b[96mdecidir sobre la petición de aplazamiento de la audiencia sobre medidas provisionales\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 567, 512, 579]]\u001b[0m\n",
      "\u001b[96mformuladas por el Gobierno de Guatemala.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 578, 300, 590]]\u001b[0m\n",
      "\u001b[96mPOR TANTO:\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 602, 159, 614]]\u001b[0m\n",
      "\u001b[96mLA CORTE INTERAMERICANA DE DERECHOS HUMANOS,\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 626, 393, 638]]\u001b[0m\n",
      "\u001b[96mRESUELVE:\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 649, 151, 661]]\u001b[0m\n",
      "\u001b[96m1.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 673, 103, 685]]\u001b[0m\n",
      "\u001b[96mConvocar a una audiencia pública para el 30 de julio de 1991 a las 15:00\u001b[0m \u001b[93m[Page: 1 || Coordinates: [122, 673, 513, 685]]\u001b[0m\n",
      "\u001b[96mhoras con el objeto de conocer los puntos de vista del Gobierno de Guatemala y de la\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 685, 512, 697]]\u001b[0m\n",
      "\u001b[96mComisión sobre la solicitud de prórroga formulada por el primero.\u001b[0m \u001b[93m[Page: 1 || Coordinates: [88, 697, 412, 709]]\u001b[0m\n",
      "\u001b[96m2\u001b[0m \u001b[93m[Page: 2 || Coordinates: [294, 71, 300, 83]]\u001b[0m\n",
      "\u001b[96m2.\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 106, 101, 118]]\u001b[0m\n",
      "\u001b[96mConocer también, en dicha audiencia pública, de las medidas que, en atención\u001b[0m \u001b[93m[Page: 2 || Coordinates: [122, 106, 511, 118]]\u001b[0m\n",
      "\u001b[96ma la resolución del Presidente del 15 de julio del presente año, ha tomado el\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 118, 515, 130]]\u001b[0m\n",
      "\u001b[96mGobierno de Guatemala.\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 129, 211, 141]]\u001b[0m\n",
      "\u001b[96m3.\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 153, 103, 165]]\u001b[0m\n",
      "\u001b[96mReservarse el derecho de convocar a una audiencia pública para resolver la\u001b[0m \u001b[93m[Page: 2 || Coordinates: [122, 153, 513, 165]]\u001b[0m\n",
      "\u001b[96mpetición de la Comisión sobre medidas provisionales respecto de Guatemala.\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 165, 467, 177]]\u001b[0m\n",
      "\u001b[96mHéctor Fix-Zamudio\u001b[0m \u001b[93m[Page: 2 || Coordinates: [249, 200, 349, 212]]\u001b[0m\n",
      "\u001b[96mPresidente\u001b[0m \u001b[93m[Page: 2 || Coordinates: [272, 212, 327, 224]]\u001b[0m\n",
      "\u001b[96mOrlando\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 248, 161, 260]]\u001b[0m\n",
      "\u001b[96mTovar\u001b[0m \u001b[93m[Page: 2 || Coordinates: [129, 248, 191, 260]]\u001b[0m\n",
      "\u001b[96mTamayo\u001b[0m \u001b[93m[Page: 2 || Coordinates: [161, 248, 234, 260]]\u001b[0m\n",
      "\u001b[96mThomas\u001b[0m \u001b[93m[Page: 2 || Coordinates: [225, 248, 436, 260]]\u001b[0m\n",
      "\u001b[96mBuergenthal\u001b[0m \u001b[93m[Page: 2 || Coordinates: [405, 248, 499, 260]]\u001b[0m\n",
      "\u001b[96mRafael Nieto Navia\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 283, 195, 295]]\u001b[0m\n",
      "\u001b[96mPolicarpo Callejas Bonilla\u001b[0m \u001b[93m[Page: 2 || Coordinates: [329, 283, 481, 295]]\u001b[0m\n",
      "\u001b[96mSonia\u001b[0m \u001b[93m[Page: 2 || Coordinates: [88, 318, 150, 330]]\u001b[0m\n",
      "\u001b[96mPicado\u001b[0m \u001b[93m[Page: 2 || Coordinates: [118, 318, 184, 330]]\u001b[0m\n",
      "\u001b[96mSotela\u001b[0m \u001b[93m[Page: 2 || Coordinates: [153, 318, 218, 330]]\u001b[0m\n",
      "\u001b[96mJulio\u001b[0m \u001b[93m[Page: 2 || Coordinates: [191, 318, 419, 330]]\u001b[0m\n",
      "\u001b[96mA.\u001b[0m \u001b[93m[Page: 2 || Coordinates: [388, 318, 433, 330]]\u001b[0m\n",
      "\u001b[96mBarberis\u001b[0m \u001b[93m[Page: 2 || Coordinates: [402, 318, 477, 330]]\u001b[0m\n",
      "\u001b[96mManuel E. Ventura Robles\u001b[0m \u001b[93m[Page: 2 || Coordinates: [235, 354, 364, 366]]\u001b[0m\n",
      "\u001b[96mSecretario\u001b[0m \u001b[93m[Page: 2 || Coordinates: [273, 366, 326, 378]]\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "for page in pdf_features.pages:\n",
    "    for token in page.tokens:\n",
    "        coordinates = [token.bounding_box.left, token.bounding_box.top, token.bounding_box.right, token.bounding_box.bottom]\n",
    "        print(f\"\\033[96m{token.content}\\033[0m \\033[93m[Page: {page.page_number} || Coordinates: {coordinates}]\\033[0m\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4576ff4d-92fc-4e19-a947-ebfb3fd01060",
   "metadata": {},
   "source": [
    "## Fine-Tuning Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01826a89-25c9-4385-a1e6-b65c0edbd0c6",
   "metadata": {},
   "source": [
    "Now that we have some overview about the `PdfFeatures` module, we can now start fine-tuning process."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "586eba43-9138-4eff-a3fa-24553de04e82",
   "metadata": {},
   "source": [
    "In the `pdf-document-layout-analysis` service, there are two LightGBM (i.e. fast) models.\n",
    "\n",
    "- The first model is used to determine the types of tokens. We call it `token_type_model`.\n",
    "- The second model is used to identify the segments to which the tokens belong. We call this model `paragraph_extraction_model`.\n",
    "\n",
    "The second model uses the predictions from the first model's output (predicted token types) as part of its features. So, let's start by fine-tuning the token type model."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c326ccb1-a36b-40f9-b7e2-e83ba3c0e12b",
   "metadata": {},
   "source": [
    "### Fine-Tuning Token Type Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b3638eb-c512-4bd5-97f4-4df3ae984978",
   "metadata": {},
   "source": [
    "#### Loading Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab35b27c-8464-470c-9ef1-a9aef8945f6a",
   "metadata": {},
   "source": [
    "To properly train a token type model, you should have a list of PdfFeatures items where the `token_type` attribute of their tokens is set correctly, as this attribute will be used as the label.\n",
    "\n",
    "To see what labels are going to be used in the model, you can check `pdf_token_type_labels/TokenType.py`. As default, we are using the labels of [DocLayNet](https://github.com/DS4SD/DocLayNet) dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2ab3093c-6e67-4505-bac3-b7db73ef5372",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_pdf_features_labels() -> PdfFeatures:\n",
    "    # Assuming that you are loading your own labels in this part.\n",
    "    # I'm just going to put a list with a single file for demonstration.\n",
    "    pdf_features: PdfFeatures = PdfFeatures.from_pdf_path(\"test_pdfs/regular.pdf\")\n",
    "    labeled_pdf_features_list: list[PdfFeatures] = [pdf_features]\n",
    "    return labeled_pdf_features_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32db8aee-9d2c-45bf-b7af-ac6249081f32",
   "metadata": {},
   "source": [
    "If you want to use it, there is also an optional method inside `PdfFeatures` to load the labeled data. It can be used as follows:\n",
    "\n",
    "\n",
    "    pdf_features: PdfFeatures = PdfFeatures.from_labeled_data(pdf_labeled_data_root_path, dataset_name, pdf_name)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fda0c166-ac25-4084-974a-c73f1cb06f18",
   "metadata": {},
   "source": [
    "But to be able to use it, you should have this file structure in somewhere else for your labeled data:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5acf2beb-f7a2-4e12-8f11-4bffff7efa74",
   "metadata": {},
   "source": [
    "```\n",
    ".\n",
    "└── pdf-labeled-data\n",
    "    ├── labeled_data\n",
    "    │   └── token_type\n",
    "    │       ├── train_data\n",
    "    │       │   ├── document1\n",
    "    │       │   │   └── labels.json\n",
    "    │       │   ├── document2\n",
    "    │       │   │   └── labels.json\n",
    "    │       │   └── document3\n",
    "    │       │       └── labels.json\n",
    "    │       └── test_data\n",
    "    │           └── document4\n",
    "    │               └── labels.json\n",
    "    └── pdfs\n",
    "        ├── document1\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        ├── document2\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        ├── document3\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        └── document4\n",
    "            ├── document.pdf\n",
    "            └── etree.xml\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c50cbae-9841-4289-9097-7357a0c724a7",
   "metadata": {},
   "source": [
    "Some details about this structure:\n",
    "\n",
    "- `pdf_labeled_data_root_path` parameter in `PdfFeatures.from_labeled_data` is the path to `pdf-labeled-data`\n",
    "- Directory names `labeled_data`, `token_type` and `pdfs` should not be something else, they have to be named as like this.\n",
    "- File names `labels.json`, `document.pdf` and `etree.xml` should not be something else, they have to be named as like this.\n",
    "- `token_type` directory is where your datasets are located, so their names can be anything.\n",
    "- Inside the dataset directories, there will be directories for each document. Each directory should be named according to its document name.\n",
    "- Inside the directory with the document name, there will be `labels.json` file (the structure of it will be shown below).\n",
    "- Inside `pdfs` directory, there will be a directory for each document. And inside every directory, there will be `document.pdf` and `etree.xml`.\n",
    "- `document.pdf` is the document itself and `etree.xml` is the output of Poppler.\n",
    "- Let's say we have two documents `example.pdf` and `example2.pdf`. Then, the structure should be like:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19abde59-7ba5-4e65-8ce7-6bb7fb2202d5",
   "metadata": {},
   "source": [
    "```\n",
    ".\n",
    "└── pdf-labeled-data\n",
    "    ├── labeled_data\n",
    "    │   └── token_type\n",
    "    │       ├── train_data\n",
    "    │       │   └── example1\n",
    "    │       │       └── labels.json\n",
    "    │       └── test_data\n",
    "    │           └── example2\n",
    "    │               └── labels.json\n",
    "    └── pdfs\n",
    "        ├── example1\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        └── example2\n",
    "            ├── document.pdf\n",
    "            └── etree.xml\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5379e82a-9fa7-4fea-9d6b-a11e672707bc",
   "metadata": {},
   "source": [
    "For `labels.json` files, they should have this structure:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5b7f4fb-7052-4e8c-856a-6b1d83e5ece4",
   "metadata": {},
   "source": [
    "```\n",
    "{\n",
    "    \"pages\": [\n",
    "        {\n",
    "            \"number\": 1,\n",
    "            \"labels\": [\n",
    "                {\n",
    "                    \"top\": 106,\n",
    "                    \"left\": 244,\n",
    "                    \"width\": 111,\n",
    "                    \"height\": 12,\n",
    "                    \"label_type\": 8,\n",
    "                    \"metadata\": \"\"\n",
    "                },\n",
    "                {\n",
    "                    \"top\": 118,\n",
    "                    \"left\": 157,\n",
    "                    \"width\": 284,\n",
    "                    \"height\": 12,\n",
    "                    \"label_type\": 8,\n",
    "                    \"metadata\": \"\"\n",
    "                }\n",
    "            ]\n",
    "        },\n",
    "        {\n",
    "            \"number\": 2,\n",
    "            \"labels\": [\n",
    "                {\n",
    "                    \"top\": 71,\n",
    "                    \"left\": 294,\n",
    "                    \"width\": 6,\n",
    "                    \"height\": 12,\n",
    "                    \"label_type\": 6,\n",
    "                    \"metadata\": \"\"\n",
    "                },\n",
    "                {\n",
    "                    \"top\": 106,\n",
    "                    \"left\": 88,\n",
    "                    \"width\": 13,\n",
    "                    \"height\": 12,\n",
    "                    \"label_type\": 6,\n",
    "                    \"metadata\": \"\"\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6808202-892d-43e0-9e7a-73ebc347901f",
   "metadata": {},
   "source": [
    "Here you see a list of labels for each page. Each label includes information about the coordinates `top`, `left`, `width`, `height` and their types `label_type` for each token."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b31a859-7867-4bd0-be13-7ae4ff4c8a61",
   "metadata": {},
   "source": [
    "The value for `label_type` here is the index number in the `TokenType` class.  For example, if you check `pdf_token_type_labels/TokenType.py`, the `label_type` value for `FORMULA` is `0`, for `TEXT` it's `6` and so on."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2778fd0b-5351-4c83-a15a-ecf8aac91397",
   "metadata": {},
   "source": [
    "Using this information, you can load your data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8112645a-816d-4579-b6e9-14b505703fc9",
   "metadata": {},
   "source": [
    "#### Fine-Tuning the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b96c7988-cd1e-492a-9990-84db9f7111d2",
   "metadata": {},
   "source": [
    "After we finished with loading data, we can now continue our process. This part will be simpler."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c40e426-af77-47fc-a82c-77b5ca4fddeb",
   "metadata": {},
   "source": [
    "First, to be able to use our trained token type model, you should download it from our huggingface repo. To do that, you can simply run `download_models.py` and it will download the both models."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e234e69-7e50-4ffe-a31b-2dc8248a676f",
   "metadata": {},
   "source": [
    "If you want to download it manually, you can use this link: https://huggingface.co/HURIDOCS/pdf-document-layout-analysis/tree/main  \n",
    "\n",
    "After downloading it, place it into `models` directory. The path should be as follows:  \n",
    "`~/pdf-document-layout-analysis/models/token_type_lightgbm.model`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "472072a6-a02c-4b75-bbc0-f13bb7e357d2",
   "metadata": {},
   "source": [
    "Now, training part."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb6e716b-b742-4186-9e1a-ac5ecea708ac",
   "metadata": {},
   "source": [
    "If you want to train your own model from scratch, you can use this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2c8a9b3-6180-41f2-bb82-bea892a61f5e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pdf_tokens_type_trainer.ModelConfiguration import ModelConfiguration\n",
    "from pdf_tokens_type_trainer.TokenTypeTrainer import TokenTypeTrainer\n",
    "\n",
    "def train_token_type_model():\n",
    "    model_configuration = ModelConfiguration()\n",
    "    labeled_pdf_features_list: list[PdfFeatures] = get_pdf_features_labels()\n",
    "    trainer = TokenTypeTrainer(labeled_pdf_features_list, model_configuration)\n",
    "    train_labels = [token.token_type.get_index() for token in trainer.loop_tokens()]\n",
    "    trainer.train(\"models/token_type_example_model.model\", train_labels)    \n",
    "\n",
    "train_token_type_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb6ae549-4f52-45b0-853a-6414ca8b4af3",
   "metadata": {},
   "source": [
    "Don't forget to check what's inside the `model_configuration`. You might need to tune the hyperparameters."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf3f6a6c-cba7-43c4-9f72-85cbe447cb6e",
   "metadata": {},
   "source": [
    "If you want to use our trained models as base and refit with your own data, you can use this function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "29dbaba4-d3d6-4985-be44-df872fe9b5d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def refit_token_type_model():\n",
    "    model_configuration = ModelConfiguration()\n",
    "    model_configuration.resume_training = True\n",
    "    labeled_pdf_features_list: list[PdfFeatures] = get_pdf_features_labels()\n",
    "    trainer = TokenTypeTrainer(labeled_pdf_features_list, model_configuration)\n",
    "    train_labels = [token.token_type.get_index() for token in trainer.loop_tokens()]\n",
    "    trainer.train(\"models/token_type_lightgbm.model\", train_labels)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a82f6f6-cec9-48bc-9c64-b09aa65d2754",
   "metadata": {},
   "source": [
    "Running this function will refit the same model with your data. Depending on your situation, it may or may not help you."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b95cd2cd-0d41-4518-8576-b1a0d2adc21b",
   "metadata": {},
   "source": [
    "If it does not help, you can try to check other fine-tuning strategies in LightGBM. \n",
    "\n",
    "In that case, all you need to do is changing this part in `pdf_tokens_type_trainer/PdfTrainer.py` (lines 47-49):\n",
    "\n",
    "```\n",
    "        if self.model_configuration.resume_training and exists(model_path):\n",
    "            model = lgb.Booster(model_file=model_path)\n",
    "            gbm = model.refit(x_train, labels)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67948603-80e6-4b42-9ba1-78868fd9f946",
   "metadata": {},
   "source": [
    "To make predictions with the trained model, you can use this function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2e7cd129-874e-415d-9855-401d8c5d0136",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Page-1\n",
      "Page-2\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 87.79it/s]\n",
      "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 466.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[96mRESOLUCIÓN DE LA\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mCORTE INTERAMERICANA DE DERECHOS HUMANOS\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mDEL 29 DE JULIO DE 1991\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mMEDIDAS PROVISIONALES SOLICITADAS POR LA COMISIÓN\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mINTERAMERICANA DE DERECHOS HUMANOS\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mRESPECTO DE GUATEMALA\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mCASO CHUNIMA\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mLA CORTE INTERAMERICANA DE DERECHOS HUMANOS,\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96mVISTOS:\u001b[0m \u001b[93m[TokenType.TITLE]\u001b[0m\n",
      "\u001b[96m1.\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mLa resolución del Presidente de la Corte Interamericana de Derechos Humanos\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mde 15 de julio de 1991, sobre medidas provisionales solicitadas por la Comisión\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mInteramericana de Derechos Humanos respecto de Guatemala;\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96m2.\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mLa convocatoria a una audiencia pública para el día 29 de julio de 1991 a las\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96m3:00 p.m., contenida en la resolución citada;\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96m3.\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mLos escritos de fechas 24 y 26 de este mes de julio presentados por el\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mGobierno de Guatemala en los cuales informa que, en atención a la resolución del\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n",
      "\u001b[96mPresidente, ha tomado medidas dirigidas a la protección de las personas\u001b[0m \u001b[93m[TokenType.TEXT]\u001b[0m\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def get_predictions():\n",
    "    model_configuration = ModelConfiguration()\n",
    "    pdf_features: PdfFeatures = PdfFeatures.from_pdf_path(\"test_pdfs/regular.pdf\")\n",
    "    trainer = TokenTypeTrainer([pdf_features], model_configuration)\n",
    "    trainer.set_token_types()\n",
    "    for token in pdf_features.pages[0].tokens[:20]:\n",
    "        print(f\"\\033[96m{token.content}\\033[0m \\033[93m[{token.token_type}]\\033[0m\")\n",
    "\n",
    "get_predictions()    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37b6b980-deaf-4ba4-baf0-7bf137af63a7",
   "metadata": {},
   "source": [
    "### Fine-Tuning Paragraph Extraction Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1389cf49-c163-4f90-ab0c-9606756b8ef9",
   "metadata": {},
   "source": [
    "#### Loading Data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1d4cf8c-65d2-4496-adcf-ab73acc5000f",
   "metadata": {},
   "source": [
    "The second model in our pipeline is the paragraph extraction model. After finding the type of each token, now, we are going to \"group\" the tokens, which means, we are going to find each token's segment."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e747aa-9b19-4e8d-acbb-f8d221dfe006",
   "metadata": {},
   "source": [
    "We are going to explain the process but for this part, we highly recommend you to place your labeled data as in this following file structure and use the already existing methods. Otherwise, it can be a bit more harder for you to use our modules:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3af70a1-404e-4bac-a366-f7962636b1eb",
   "metadata": {},
   "source": [
    "```\n",
    ".\n",
    "└── pdf-labeled-data\n",
    "    ├── labeled_data\n",
    "    │   ├── token_type\n",
    "    │   │   ├── train_data\n",
    "    │   │   │   ├── example_document1\n",
    "    │   │   │   │   └── labels.json\n",
    "    │   │   │   ├── example_document2\n",
    "    │   │   │   │   └── labels.json\n",
    "    │   │   │   └── example_document3\n",
    "    │   │   │       └── labels.json\n",
    "    │   │   └── test_data\n",
    "    │   │       └── example_document4\n",
    "    │   │           └── labels.json\n",
    "    │   └── paragraph_extraction\n",
    "    │       ├── train_data\n",
    "    │       │   ├── example_document1\n",
    "    │       │   │   └── labels.json\n",
    "    │       │   ├── example_document2\n",
    "    │       │   │   └── labels.json\n",
    "    │       │   └── example_document3\n",
    "    │       │       └── labels.json\n",
    "    │       └── test_data\n",
    "    │           └── example_document4\n",
    "    │               └── labels.json\n",
    "    └── pdfs\n",
    "        ├── example_document1\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        ├── example_document2\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        ├── example_document3\n",
    "        │   ├── document.pdf\n",
    "        │   └── etree.xml\n",
    "        └── example_document4\n",
    "            ├── document.pdf\n",
    "            └── etree.xml\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dc0c106-7b22-42e3-969f-d52ecddae3ae",
   "metadata": {},
   "source": [
    "Some details about this structure:\n",
    "\n",
    "- Every detail in the token type labels file structure applies for this structure too.\n",
    "- `paragraph_extraction` directory is where your paragraph extraction datasets are located, its name should not be something else.\n",
    "- `token_type` labels are also shown in the structure because token types are used as a feature in the paragraph extraction model. If you do not have it, it will not break the pipeline and still train the model but the token_type feature for every token will be `TokenType.TEXT` in paragraph extractor model's features.\n",
    "- If you do not have `token_type` labels, another option is, after loading the data, you can predict the token types with the token type model (will be shown below)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e47af60-b238-43c3-9ca7-58ee8041f4f7",
   "metadata": {},
   "source": [
    "For labels.json files, they should have this structure:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d5b2376-d983-4c49-8130-b94368782828",
   "metadata": {},
   "source": [
    "```\n",
    "{\n",
    "    \"pages\": [\n",
    "        {\n",
    "            \"number\": 1,\n",
    "            \"labels\": [\n",
    "                {\n",
    "                    \"top\": 86,\n",
    "                    \"left\": 162,\n",
    "                    \"width\": 292,\n",
    "                    \"height\": 24,\n",
    "                    \"label_type\": 0\n",
    "                },\n",
    "                {\n",
    "                    \"top\": 122,\n",
    "                    \"left\": 221,\n",
    "                    \"width\": 174,\n",
    "                    \"height\": 12,\n",
    "                    \"label_type\": 0\n",
    "                }\n",
    "            ]\n",
    "        },\n",
    "        {\n",
    "            \"number\": 2,\n",
    "            \"labels\": [\n",
    "                {\n",
    "                    \"top\": 36,\n",
    "                    \"left\": 296,\n",
    "                    \"width\": 22,\n",
    "                    \"height\": 13,\n",
    "                    \"label_type\": 0\n",
    "                },\n",
    "                {\n",
    "                    \"top\": 72,\n",
    "                    \"left\": 71,\n",
    "                    \"width\": 473,\n",
    "                    \"height\": 49,\n",
    "                    \"label_type\": 0\n",
    "                }\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db009241-7d04-416a-9b75-d947626b021e",
   "metadata": {},
   "source": [
    "Here you see a list of labels for each page. Each label includes information about the coordinates `top`, `left`, `width`, `height` for each segment/paragraph. So, this time the coordinates belongs to the segments, not to the tokens.\n",
    "\n",
    "As \"label_type\", it should be always 0 since there is only one type \"paragraph\" (don't get confused with this part, it's not important, just put 0 and go on).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1972189b-c70b-436d-9830-56adc354b777",
   "metadata": {},
   "source": [
    "Using this information, you can load your data like this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "c924b60b-d70c-43b4-814f-669adf8ddf23",
   "metadata": {},
   "outputs": [],
   "source": [
    "from os.path import join\n",
    "from paragraph_extraction_trainer.PdfParagraphTokens import PdfParagraphTokens\n",
    "\n",
    "\n",
    "def load_paragraph_extraction_labels():\n",
    "\t\n",
    "\tpdf_labeled_data_root_path = \"path/to/pdf/labeled/data\"\n",
    "\tdatasets_path = join(pdf_labeled_data_root_path, \"paragraph_extraction\")\n",
    "\tlabeled_data: list[PdfParagraphTokens] = []\n",
    "\t\n",
    "\tfor dataset in listdir(join(datasets_path)):\n",
    "\t\tif \"train\" not in dataset:\n",
    "\t\t\tcontinue\n",
    "\t\tpdf_paragraph_tokens: PdfParagraphTokens = PdfParagraphTokens.from_labeled_data(pdf_labeled_data_root_path, dataset, pdf_name)\n",
    "\t\tlabeled_data.append(pdf_paragraph_tokens)\n",
    "\t\n",
    "\treturn labeled_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6c07ba4-334e-4ff3-8e2f-b2f684f053c9",
   "metadata": {},
   "source": [
    "In case you do not have `token_type` labels and want to find the types with the `token_type_model`, you can use this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "63eb5339-a0e3-4d6f-9c7e-e4a9ec3b4f43",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pdf_tokens_type_trainer.ModelConfiguration import ModelConfiguration\n",
    "from pdf_tokens_type_trainer.TokenTypeTrainer import TokenTypeTrainer\n",
    "\n",
    "def load_paragraph_extraction_labels():\n",
    "\n",
    "    pdf_labeled_data_root_path = \"path/to/pdf/labeled/data\"\n",
    "    datasets_path = join(pdf_labeled_data_root_path, \"paragraph_extraction\")\n",
    "    labeled_pdf_paragraph_tokens_list: list[PdfParagraphTokens] = []\n",
    "    \n",
    "    for dataset in listdir(join(datasets_path)):\n",
    "        if \"train\" not in dataset:\n",
    "            continue\n",
    "        pdf_paragraph_tokens: PdfParagraphTokens = PdfParagraphTokens.from_labeled_data(pdf_labeled_data_root_path, dataset, pdf_name)\n",
    "        labeled_pdf_paragraph_tokens_list.append(pdf_paragraph_tokens)\n",
    "    \n",
    "    \n",
    "        token_type_model_configuration = ModelConfiguration()\n",
    "        labeled_pdf_features_list = [pdf_paragraph_tokens.pdf_features for pdf_paragraph_tokens in labeled_pdf_paragraph_tokens_list]\n",
    "    trainer = TokenTypeTrainer(labeled_pdf_features_list, model_configuration)\n",
    "    \n",
    "    \n",
    "    return labeled_pdf_paragraph_tokens_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41b6bb64-92a2-4b75-95f9-a934c104b7c0",
   "metadata": {},
   "source": [
    "#### Fine-Tuning the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a397319e-c079-46c4-84a9-3c487f5f4a14",
   "metadata": {},
   "source": [
    "Again, to be able to use our trained paragraph extraction model, you should download it from our huggingface repo. You can just run `download_models.py` and both models are going to be downloaded."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd38ced0-2925-4fe5-98ec-b633a19b5ce3",
   "metadata": {},
   "source": [
    "If you want to download it manually, you can use this link: https://huggingface.co/HURIDOCS/pdf-document-layout-analysis/tree/main\n",
    "\n",
    "After downloading it, place it into `models` directory. The path should be as follows:  \n",
    "`~/pdf-document-layout-analysis/models/paragraph_extraction_lightgbm.model`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "972623c2-8dcb-464b-956f-041c1450e6b6",
   "metadata": {},
   "source": [
    "To train the paragraph extraction model from scratch:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "60b22be7-35d0-4c34-891e-c67d25942c72",
   "metadata": {},
   "outputs": [],
   "source": [
    "from paragraph_extraction_trainer.model_configuration import MODEL_CONFIGURATION\n",
    "\n",
    "\n",
    "def loop_pdf_paragraph_tokens(pdf_paragraph_tokens_list: list[PdfParagraphTokens]):\n",
    "    for pdf_paragraph_tokens in pdf_paragraph_tokens_list:\n",
    "        for page in pdf_paragraph_tokens.pdf_features.pages:\n",
    "            if not page.tokens:\n",
    "                continue\n",
    "            for token, next_token in zip(page.tokens, page.tokens[1:]):\n",
    "                yield pdf_paragraph_tokens, token, next_token\n",
    "            yield pdf_paragraph_tokens, page.tokens[-1], page.tokens[-1]\n",
    "\n",
    "\n",
    "def train_paragraph_extraction_model():\n",
    "    labeled_pdf_paragraph_tokens_list: list[PdfParagraphTokens] = load_paragraph_extraction_labels()\n",
    "    labeled_pdf_features_list = [pdf_paragraph_tokens.pdf_features for pdf_paragraph_tokens in labeled_pdf_paragraph_tokens_list]\n",
    "    trainer = ParagraphExtractorTrainer(labeled_pdf_features_list, MODEL_CONFIGURATION)\n",
    "    \n",
    "    train_labels = []\n",
    "    for pdf_paragraph_tokens, token, next_token in loop_pdf_paragraph_tokens([pdf_paragraph_tokens]):\n",
    "        train_labels.append(pdf_paragraph_tokens.check_same_paragraph(token, next_token))\n",
    "\n",
    "    trainer.train(\"models/paragraph_extraction_example_model.model\", train_labels)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c3f2f2a-fdb9-4661-bdd3-445e212dc583",
   "metadata": {},
   "source": [
    "And to refit the model with your own data, all you need to do is setting `resume_training` configuration to `True`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "5a652ca1-b9c7-4731-ba8b-aa98cd0d11a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def refit_paragraph_extraction_model():\n",
    "    labeled_pdf_paragraph_tokens_list: list[PdfParagraphTokens] = load_paragraph_extraction_labels()\n",
    "    labeled_pdf_features_list = [pdf_paragraph_tokens.pdf_features for pdf_paragraph_tokens in labeled_pdf_paragraph_tokens_list]\n",
    "    MODEL_CONFIGURATION.resume_training = True\n",
    "    trainer = ParagraphExtractorTrainer(labeled_pdf_features_list, MODEL_CONFIGURATION)\n",
    "    \n",
    "    train_labels = []\n",
    "    for pdf_paragraph_tokens, token, next_token in loop_pdf_paragraph_tokens([pdf_paragraph_tokens]):\n",
    "        train_labels.append(pdf_paragraph_tokens.check_same_paragraph(token, next_token))\n",
    "\n",
    "    trainer.train(\"models/paragraph_extraction_example_model.model\", train_labels)    "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "602b3174-02de-4035-bf2d-22df20ae4883",
   "metadata": {},
   "source": [
    "<font color='red'>[IMPORTANT]</font> If you want to use your own trained models in pdf-document-layout-analysis service, make sure their names are `token_type_lightgbm.model` and `paragraph_extraction_lightgbm.model` and are placed in `models` directory."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ca5d8ef-7455-4723-af4e-d8c49096251f",
   "metadata": {},
   "source": [
    "After finishing training, you can get the predictions of the model like shown in below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73a39b6a-f712-4f9c-88a1-5b8e7aee410a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pdf_tokens_type_trainer.ModelConfiguration import ModelConfiguration\n",
    "from fast_trainer.model_configuration import MODEL_CONFIGURATION as PARAGRAPH_EXTRACTION_CONFIGURATION\n",
    "from fast_trainer.PdfSegment import PdfSegment\n",
    "from fast_trainer.ParagraphExtractorTrainer import ParagraphExtractorTrainer\n",
    "\n",
    "def get_predictions():\n",
    "    pdf_features: PdfFeatures = PdfFeatures.from_pdf_path(\"test_pdfs/regular.pdf\")\n",
    "    # First, use token type model to find and set the types.\n",
    "    token_type_trainer = TokenTypeTrainer([pdf_features], ModelConfiguration())\n",
    "    token_type_trainer.set_token_types(\"models/token_type_lightgbm.model\")\n",
    "    trainer = ParagraphExtractorTrainer(pdfs_features=[pdf_features], model_configuration=PARAGRAPH_EXTRACTION_CONFIGURATION)\n",
    "    segments: list[PdfSegment] = trainer.get_pdf_segments(\"models/paragraph_extraction_lightgbm.model\")\n",
    "    model_configuration = ModelConfiguration()\n",
    "    for segment in segments[:20]:\n",
    "        print(f\"\\033[96m{segment.text_content}\\033[0m \\033[93m[{segment.segment_type}]\\033[0m \\033[91m{segment.bounding_box.to_dict()}\\033[0m\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5a5ab63-7931-40e0-8f51-43e3f3ef5b32",
   "metadata": {},
   "source": [
    "Output of the `paragraph_extraction_model` is a list of `PdfSegment` items. Every item includes the information like `page_number`, `text_content`, `segment_type`, `bounding_box`, `pdf_name` for each of the segments. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2408b4f8-c0e8-4e92-a42a-eebde6583892",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
